import cv2
import time
import numpy as np
from stereo.utils.stereoRectify_process import stereo_rectify
from stereo.utils.config import Config

import threading
import queue
import sys
import os

class StereoCameraCaptureThreaded_rectified:
    def __init__(self, left_device, right_device, width=640, height=480, fps=50):
        self.left_device = left_device
        self.right_device = right_device
        self.width = width
        self.height = height
        self.fps = fps
        self.frame_queue = queue.Queue()  # 限制队列最大容量为5
        self.cap_left = self._open_camera(self.left_device)
        self.cap_right = self._open_camera(self.right_device)
        self.running = False
        self.capture_thread = None
        
        # 帧率控制相关变量
        self.target_fps = fps  # 默认目标帧率与摄像头采集帧率相同
        self.last_frame_time = 0
        self.frame_interval = 1.0 / self.target_fps if self.target_fps > 0 else 0
        
        # 捕获控制变量
        self.capture_fps = fps  # 捕获帧率，默认与摄像头帧率相同
        self.last_capture_time = 0
        self.capture_interval = 1.0 / self.capture_fps if self.capture_fps > 0 else 0
        
        # 加载相机标定参数
        self.map1x, self.map1y, self.map2x, self.map2y = stereo_rectify.map1x, stereo_rectify.map1y, stereo_rectify.map2x, stereo_rectify.map2y
        
        self.start_capture()

    def _open_camera(self, device_name):
        gst_pipeline = (
            f'v4l2src device=/dev/{device_name} ! '
            f'video/x-raw,format=NV12,width={self.width},height={self.height},framerate={self.fps}/1 ! '
            # f'videoscale ! video/x-raw,width=640,height=480 ! '
            f'appsink'
        )
        
        cap = cv2.VideoCapture(gst_pipeline, cv2.CAP_GSTREAMER)
        if not cap.isOpened():
            print(f"无法打开摄像头: {device_name}")
            return None
        # 打印实际的帧尺寸
        actual_width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
        actual_height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
        print(f"[{device_name}] 实际帧大小: {actual_width}x{actual_height}")
        return cap

    def apply_rectification_single(self, img, map1x, map1y, k=None, dist=None):
        if k is not None:
            img = cv2.undistort(img, k, dist)
        rectified = cv2.remap(img, map1x, map1y, cv2.INTER_LINEAR)
        return rectified

    def _process_frame(self, frame, is_left=True):
        try:
            frame = cv2.cvtColor(frame, cv2.COLOR_YUV2BGR_NV12)
            # frame = cv2.rotate(frame, cv2.ROTATE_180)
            
            frame = self.apply_rectification_single(
                frame, 
                self.map1x if is_left else self.map2x,
                self.map1y if is_left else self.map2y,
                # self.k_left,
                # self.dist_left
            )
            # frame = cv2.resize(frame, (640, 480))
            
            return frame
        except Exception as e:
            print(f"图像处理失败: {e}")
            return None

    def _capture_loop(self):
        # 设置线程亲和性到小核 0-3
        try:
            tid = threading.get_native_id()
            os.sched_setaffinity(tid, {0,1,2,3,4,5,6,7})
            print(f"[INFO] 采集线程绑定到小核: {sorted(os.sched_getaffinity(tid))}")
        except Exception as e:
            print(f"[WARN] 设置线程亲和失败: {e}")

        # 帧率统计相关变量
        frame_count = 0
        fps_start_time = time.time()
        fps_report_interval = 5.0  # 每5秒打印一次采集帧率

        while self.running:
            current_time = time.time()
            elapsed = current_time - self.last_capture_time

            if elapsed < self.capture_interval:
                time.sleep(self.capture_interval - elapsed)

            if self.frame_queue.full():
                time.sleep(0.01)
                continue

            ret_left, frame_left = self.cap_left.read()
            ret_right, frame_right = self.cap_right.read()

            self.last_capture_time = time.time()

            if ret_left and ret_right:
                processed_left = self._process_frame(frame_left, True)
                processed_right = self._process_frame(frame_right, False)

                if processed_left is not None and processed_right is not None:
                    try:
                        self.frame_queue.put((processed_left, processed_right, time.time()), block=False)
                        frame_count += 1
                    except queue.Full:
                        pass
            else:
                if not ret_left:
                    print("未能从左摄像头读取帧")
                if not ret_right:
                    print("未能从右摄像头读取帧")
                time.sleep(0.01)

            # 每5秒打印一次采集帧率
            if time.time() - fps_start_time >= fps_report_interval:
                duration = time.time() - fps_start_time
                actual_fps = frame_count / duration
                print(f"[采集线程] 实际采集帧率: {actual_fps:.2f} FPS")
                fps_start_time = time.time()
                frame_count = 0

    def start_capture(self):
        self.running = True
        self.capture_thread = threading.Thread(target=self._capture_loop, daemon=True)
        self.capture_thread.start()
        
        
    def set_capture_fps(self, capture_fps):
        """设置捕获帧率，用于控制捕获线程的速率"""
        if capture_fps <= 0:
            print("警告: 捕获帧率必须大于0，设置为默认值1")
            capture_fps = 1
        self.capture_fps = capture_fps
        self.capture_interval = 1.0 / self.capture_fps
        print(f"捕获帧率已设置为: {self.capture_fps} fps (帧间隔: {self.capture_interval*1000:.2f}ms)")
        return True

    def get_stereo_frames(self, block=True, timeout=1.0):
        """
        获取立体相机帧，可以控制帧率
        
        参数:
            block: 如果为True，在队列为空时会阻塞等待
            timeout: 阻塞等待的超时时间（秒）
            
        返回:
            (left_frame, right_frame) 或者在超时时返回 (None, None)
        """
            
        try:
            if block:
                frame_data = self.frame_queue.get(timeout=timeout)
            else:
                if self.frame_queue.empty():
                    return None, None
                frame_data = self.frame_queue.get_nowait()
                
            # 更新上次获取帧的时间
            self.last_frame_time = time.time()
            
            # 返回左右图像，丢弃时间戳
            return frame_data[0], frame_data[1], frame_data[2]
        except queue.Empty:
            print("获取帧超时")
            return None, None, None

    def get_queue_size(self):
        """获取当前队列中的帧数"""
        return self.frame_queue.qsize()
        
    def clear_queue(self):
        """清空帧队列"""
        while not self.frame_queue.empty():
            try:
                self.frame_queue.get_nowait()
            except queue.Empty:
                break
        return True

    def close(self):
        self.running = False
        if self.capture_thread:
            self.capture_thread.join(timeout=1.0)  # 设置超时，避免无限等待
        if self.cap_left:
            self.cap_left.release()
        if self.cap_right:
            self.cap_right.release()
        print("摄像头已释放")

    def __del__(self):
        self.close()

# 更新实例化代码
StereoCameraCapture_threaded_ins = StereoCameraCaptureThreaded_rectified("video22", "video31")
# 设置目标帧率为15fps，捕获帧率也设为15fps
StereoCameraCapture_threaded_ins.set_capture_fps(30)

if __name__ == "__main__":
    save_video = False
    target_fps = 30  # 设置目标帧率为15fps
    
    # 设置目标帧率和捕获帧率
    StereoCameraCapture_threaded_ins.set_capture_fps(target_fps)
    
    if save_video:
        # 创建视频写入器
        fourcc = cv2.VideoWriter_fourcc(*'XVID')
        out = cv2.VideoWriter('stereo_output.avi', fourcc, target_fps, (640, 480*2))  # 使用目标帧率
        
        try:
            frame_count = 0
            start_time = time.time()
            
            while True:
                frame_left, frame_right = StereoCameraCapture_threaded_ins.get_stereo_frames()
                
                if frame_left is None or frame_right is None:
                    print("未能获取有效帧")
                    continue
                
                # 上下组合左右图像
                combined_frame = np.vstack((frame_left, frame_right))
                
                # 保存视频
                out.write(combined_frame)
                
                # 计算实际帧率
                frame_count += 1
                elapsed = time.time() - start_time
                if elapsed >= 5.0:  # 每5秒显示一次帧率
                    actual_fps = frame_count / elapsed
                    print(f"实际帧率: {actual_fps:.2f} fps (目标: {target_fps} fps)")
                    print(f"队列大小: {StereoCameraCapture_threaded_ins.get_queue_size()}")
                    frame_count = 0
                    start_time = time.time()
                
        except KeyboardInterrupt:
            print("程序被用户中断")
        finally:
            # 释放资源
            out.release()
            cv2.destroyAllWindows()
            StereoCameraCapture_threaded_ins.close()
    else:
        try:
            frame_count = 0
            start_time = time.time()
            
            while True:
                frame_left, frame_right, frame_time = StereoCameraCapture_threaded_ins.get_stereo_frames()
                
                if frame_left is None or frame_right is None:
                    continue
                
                # 计算实际帧率
                frame_count += 1
                elapsed = time.time() - start_time
                # if elapsed >= 5.0:  # 每5秒显示一次帧率
                #     actual_fps = frame_count / elapsed
                #     print(f"实际帧率: {actual_fps:.2f} fps (目标: {target_fps} fps)")
                #     print(f"队列大小: {StereoCameraCapture_threaded_ins.get_queue_size()}")
                #     frame_count = 0
                #     start_time = time.time()
                
        except KeyboardInterrupt:
            print("程序被用户中断")
        finally:
            StereoCameraCapture_threaded_ins.close()

